from lazyllm import Document, Retriever
import lazyllm
import os
import datetime

# 获取一周前的时间
def get_one_week_ago():
    return datetime.datetime.now() - datetime.timedelta(weeks=1)

# 过滤一周内修改的文件
def filter_recent_files(directory_path):
    one_week_ago = get_one_week_ago()
    recent_files = []
    
    # 检查目录是否存在
    if not os.path.exists(directory_path):
        return recent_files
    
    for filename in os.listdir(directory_path):
        file_path = os.path.join(directory_path, filename)
        if os.path.isfile(file_path):
            try:
                # 获取文件修改时间
                mod_time = datetime.datetime.fromtimestamp(os.path.getmtime(file_path))
                # 判断是否在一周内
                if mod_time > one_week_ago:
                    recent_files.append(file_path)
            except Exception:
                # 忽略无法获取修改时间的文件
                pass
    
    return recent_files

# 智能总结函数
def smart_summary(query, doc_path, use_time_filter=False):
    # 检查查询是否涉及最近一周的文档
    if "最近一周" in query or "近一周" in query or "一周" in query:
        use_time_filter = True
    
    # 检查文档路径是否存在
    if not os.path.exists(doc_path):
        return "文档目录不存在，请检查路径设置。"
    
    try:
        # 根据时间过滤选项决定加载哪些文件
        if use_time_filter:
            # 只加载一周内修改的文件
            file_paths = filter_recent_files(doc_path)
            if not file_paths:
                return "近一周内没有修改的文档。"
            
            # 创建文档对象，使用特定文件列表
            documents = lazyllm.Document(dataset_path=file_paths)
            time_filter_message = "（已筛选近一周内修改的文档）"
        else:
            # 加载所有文档
            documents = lazyllm.Document(dataset_path=doc_path)
            time_filter_message = ""
        
        # 创建检索器
        retriever = lazyllm.Retriever(doc=documents, group_name="CoarseChunk", similarity="bm25_chinese", topk=3)
        
        # 创建大语言模型，使用SenseNova支持的模型
        llm = lazyllm.OnlineChatModule(source="sensenova", model='SenseNova-V6-Turbo')
        
        # 设置提示词
        prompt = 'You will act as an AI question-answering assistant and complete a dialogue task. In this task, you need to provide your answers based on the given context and questions.'
        llm.prompt(lazyllm.ChatPrompter(instruction=prompt, extra_keys=['context_str']))
        
        # 检索相关文档
        docs_node_list = retriever(query=query)
        
        # 生成回答
        context_str = "".join([node.get_content() for node in docs_node_list])
        if not context_str:
            if use_time_filter:
                return "近一周内没有找到与查询相关的文档内容。"
            else:
                return "未找到与查询相关的文档内容。"
        
        res = llm({"query": query, "context_str": context_str})
        
        # 添加时间过滤信息到结果中
        if time_filter_message and use_time_filter:
            res = time_filter_message + "\n" + res
        
        return res
    except Exception as e:
        return f"处理文档时出错: {str(e)}. 请确保文档目录中包含有效的文档文件，并且LazyLLM能够访问该目录。"

# 创建 Web 界面 - 使用ActionModule包装自定义函数
def create_web_interface():
    # 使用ActionModule包装smart_summary函数作为后端处理逻辑
    chat = lazyllm.ActionModule(smart_summary)
    
    # 创建 Web 界面，添加文档路径输入和时间过滤选项
    web = lazyllm.WebModule(
        chat, 
        port=23333,
        components={
            chat: [
                ('doc_path', 'Text', ''),  # 将默认值设为空字符串，让用户手动输入
                ('use_time_filter', 'Checkbox', False)
            ]
        }
    )
    
    return web

# 主函数
if __name__ == "__main__":
    # 启动 Web 界面
    web_interface = create_web_interface()
    web_interface.start().wait()